# -*- coding: utf-8 -*-

import io
import os
import sys

from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.prompts import PromptTemplate
from langchain_ollama import OllamaLLM

# 设置环境变量
os.environ["PYTHONIOENCODING"] = "utf-8"

# 设置控制台输出编码为 UTF-8
if sys.platform.startswith('win'):
    # Windows 系统特殊处理
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
else:
    # 非 Windows 系统
    sys.stdout.reconfigure(encoding='utf-8')
    sys.stderr.reconfigure(encoding='utf-8')


class CodeReviewChain:
    """代码审查链示例"""

    def __init__(self):
        self.llm = OllamaLLM(
            model="qwen3:8b",
            callbacks=[StreamingStdOutCallbackHandler()],
            temperature=0.7,
        )
        self.prompt = PromptTemplate.from_template(
            """你是一个专业的代码审查员。请审查以下代码并提供改进建议：
            
            代码:
            {code}
            
            请从以下几个方面进行审查：
            1. 代码质量
            2. 性能优化
            3. 安全性
            4. 最佳实践
            
            审查结果:"""
        )
        self.chain = self.prompt | self.llm

    def review_code(self, code: str) -> str:
        """审查代码"""
        return self.chain.invoke({"code": code})


class DocumentQAChain:
    """文档问答链示例"""

    def __init__(self):
        self.llm = OllamaLLM(
            model="qwen3:8b",
            callbacks=[StreamingStdOutCallbackHandler()],
            temperature=0.7,
        )
        self.prompt = PromptTemplate.from_template(
            """基于以下文档内容回答问题：
            
            文档内容:
            {document}
            
            问题: {question}
            
            请提供详细的回答:"""
        )
        self.chain = self.prompt | self.llm

    def answer_question(self, document: str, question: str) -> str:
        """回答问题"""
        return self.chain.invoke({"document": document, "question": question})


def main():
    """运行链式调用示例"""
    print("\n=== 链式调用示例 ===")

    # 代码审查示例
    print("\n--- 代码审查示例 ---")
    code_review = CodeReviewChain()
    sample_code = """
    def calculate_sum(numbers):
        sum = 0
        for i in range(len(numbers)):
            sum += numbers[i]
        return sum
    """
    print("待审查代码:")
    print(sample_code)
    print("\n审查结果:")
    print(code_review.review_code(sample_code))

    # 文档问答示例
    print("\n--- 文档问答示例 ---")
    doc_qa = DocumentQAChain()
    document = """
    LangChain 是一个用于开发由语言模型驱动的应用程序的框架。
    它提供了许多工具和抽象，使得构建复杂的应用程序变得更加容易。
    主要功能包括：链式调用、记忆功能、Agent 等。
    """
    question = "LangChain 的主要功能有哪些？"
    print(f"问题: {question}")
    print("\n回答:")
    print(doc_qa.answer_question(document, question))


if __name__ == "__main__":
    main()
